use std::{sync::Arc, time::Duration};

use kameo::{
    actor::{
        supervision::{
            DefaultStrategy, ExponentialBackoffStrategy, FailureType, RateLimitedStrategy,
            SupervisionContext, SupervisionDecision, SupervisionStrategy, SupervisionStrategyBuilder,
            SupervisedActor,
        },
        Actor, ActorID, ActorRef, WeakActorRef,
    },
    error::{BoxError, PanicError},
    mailbox::unbounded::UnboundedMailbox,
    message::Message,
};
use tokio::time::sleep;

// Message that causes the actor to panic
struct Crash;

// Message to trigger link death
struct KillLinked;

// Message to print the actor's state
struct GetState;

// Reply for GetState
struct StateReply(usize);

// Implement the SupervisedActor trait for a custom actor
struct SupervisedCounterActor {
    count: usize, 
    supervision_context: SupervisionContext,
    supervision_strategy: Box<dyn SupervisionStrategy>,
}

impl SupervisedCounterActor {
    fn new(strategy: Box<dyn SupervisionStrategy>) -> Self {
        Self {
            count: 0,
            supervision_context: SupervisionContext::new(10),
            supervision_strategy: strategy,
        }
    }
    
    fn with_exponential_backoff() -> Self {
        Self::new(
            SupervisionStrategyBuilder::new()
                .with_exponential_backoff(3, Duration::from_millis(100), Duration::from_secs(2))
                .build()
        )
    }
    
    fn with_rate_limited() -> Self {
        Self::new(
            SupervisionStrategyBuilder::new()
                .with_rate_limited_restarts(3, Duration::from_secs(10))
                .build()
        )
    }
}

impl Actor for SupervisedCounterActor {
    type Mailbox = UnboundedMailbox<Self>;
    
    async fn on_start(&mut self, actor_ref: ActorRef<Self>) -> Result<(), BoxError> {
        println!("Actor started with ID: {}", actor_ref.id());
        Ok(())
    }
    
    async fn on_panic(
        &mut self, 
        actor_ref: WeakActorRef<Self>, 
        err: PanicError
    ) -> Result<Option<kameo::error::ActorStopReason>, BoxError> {
        println!("Actor panicked: {:?}", err);
        // Use our supervised panic handler
        self.handle_supervised_panic(actor_ref, err).await
    }
    
    async fn on_link_died(
        &mut self,
        actor_ref: WeakActorRef<Self>,
        id: ActorID,
        reason: kameo::error::ActorStopReason,
    ) -> Result<Option<kameo::error::ActorStopReason>, BoxError> {
        println!("Linked actor died: {} with reason: {:?}", id, reason);
        // Use our supervised link died handler
        self.handle_supervised_link_died(actor_ref, id, reason).await
    }
    
    async fn on_stop(
        &mut self,
        actor_ref: WeakActorRef<Self>,
        reason: kameo::error::ActorStopReason,
    ) -> Result<(), BoxError> {
        println!("Actor stopping with reason: {:?}", reason);
        Ok(())
    }
}

impl SupervisedActor for SupervisedCounterActor {
    fn supervision_strategy(&self) -> Box<dyn SupervisionStrategy> {
        self.supervision_strategy.clone_box()
    }
    
    fn supervision_context(&mut self) -> &mut SupervisionContext {
        &mut self.supervision_context
    }
}

impl Message<Crash> for SupervisedCounterActor {
    async fn handle(&mut self, _msg: Crash, _ctx: ActorRef<Self>) {
        println!("About to crash...");
        panic!("Actor crashed intentionally!");
    }
}

impl Message<KillLinked> for SupervisedCounterActor {
    async fn handle(&mut self, _msg: KillLinked, ctx: ActorRef<Self>) {
        println!("Killing linked actor...");
        self.count += 1;
        // Simulate a linked actor death by directly calling on_link_died
        let dummy_id = ActorID::generate();
        let reason = PanicError::new("Linked actor crashed");
        let actor_stop_reason = kameo::error::ActorStopReason::Panicked(reason);
        
        let _ = self.on_link_died(
            ctx.downgrade(), 
            dummy_id, 
            actor_stop_reason
        ).await;
    }
}

impl Message<GetState> for SupervisedCounterActor {
    type Reply = StateReply;
    
    async fn handle(&mut self, _msg: GetState, _ctx: ActorRef<Self>) -> StateReply {
        self.count += 1;
        StateReply(self.count)
    }
}

// Helper trait to enable cloning the Box<dyn SupervisionStrategy>
trait CloneableStrategy {
    fn clone_box(&self) -> Box<dyn SupervisionStrategy>;
}

impl<T> CloneableStrategy for T
where
    T: SupervisionStrategy + Clone + 'static,
{
    fn clone_box(&self) -> Box<dyn SupervisionStrategy> {
        Box::new(self.clone())
    }
}

impl Clone for Box<dyn SupervisionStrategy> {
    fn clone(&self) -> Self {
        self.as_ref().clone_box()
    }
}

// Make our strategies cloneable
impl Clone for DefaultStrategy {
    fn clone(&self) -> Self {
        Self
    }
}

impl Clone for ExponentialBackoffStrategy {
    fn clone(&self) -> Self {
        ExponentialBackoffStrategy::new(
            self.max_retries,
            self.base_duration,
            self.max_duration,
        )
    }
}

// We won't implement Clone for RateLimitedStrategy in this example
// as it contains a Box<dyn SupervisionStrategy> which complicates cloning

#[tokio::main]
async fn main() -> Result<(), BoxError> {
    // Enable logging
    tracing_subscriber::fmt()
        .with_env_filter("kameo=debug")
        .init();
    
    println!("=== Testing Exponential Backoff Strategy ===");
    test_exponential_backoff().await?;
    
    println!("\n=== Testing Rate Limited Strategy ===");
    test_rate_limited().await?;
    
    Ok(())
}

async fn test_exponential_backoff() -> Result<(), BoxError> {
    // Create an actor with exponential backoff strategy
    let actor = SupervisedCounterActor::with_exponential_backoff();
    let actor_ref = actor.spawn();
    
    // Get initial state
    let state = actor_ref.ask(GetState).await?;
    println!("Initial state: {}", state.0);
    
    // Cause the actor to panic multiple times
    for i in 0..5 {
        println!("Crash attempt {}", i + 1);
        
        // Try to crash the actor
        let _ = actor_ref.tell(Crash).send();
        
        // Wait a bit for backoff to happen
        sleep(Duration::from_secs(1)).await;
        
        // Check if the actor is still alive
        if actor_ref.is_alive() {
            println!("Actor survived crash with exponential backoff!");
            
            // Get state to see if it's still working
            match actor_ref.ask(GetState).await {
                Ok(state) => println!("Current state: {}", state.0),
                Err(e) => println!("Failed to get state: {:?}", e),
            }
        } else {
            println!("Actor died after {} crashes (expected after exceeding max retries)", i + 1);
            break;
        }
    }
    
    Ok(())
}

async fn test_rate_limited() -> Result<(), BoxError> {
    // Create an actor with rate limited strategy
    let actor = SupervisedCounterActor::with_rate_limited();
    let actor_ref = actor.spawn();
    
    // Get initial state
    let state = actor_ref.ask(GetState).await?;
    println!("Initial state: {}", state.0);
    
    // Trigger linked actor deaths
    for i in 0..5 {
        println!("Link death attempt {}", i + 1);
        
        // Kill a linked actor
        let _ = actor_ref.tell(KillLinked).send();
        
        // Small delay between attempts
        sleep(Duration::from_millis(300)).await;
        
        // Check if the actor is still alive
        if actor_ref.is_alive() {
            println!("Actor survived link death with rate limiting!");
            
            // Get state to see if it's still working
            match actor_ref.ask(GetState).await {
                Ok(state) => println!("Current state: {}", state.0),
                Err(e) => println!("Failed to get state: {:?}", e),
            }
        } else {
            println!("Actor died after {} link deaths", i + 1);
            break;
        }
    }
    
    Ok(())
} 